import heapq
import math
import random
import networkx as nx
import matplotlib.pyplot as plt
from networkx import Graph

class QKDNetwork:
    class LinkCapacity:
        def __init__(self, startNode, endNode, capacity):
            self.startNode = startNode
            self.endNode = endNode
            self.capacity = capacity

    class Node:
        def __init__(self, nodeID):
            self.nodeID = nodeID
            self.adjacentNodes = []

    def __init__(self, path = "", useNetworkx = False, networkxGraph: Graph = None):
        self.nodeList = []
        self.linkListCapacity: list[QKDNetwork.LinkCapacity] = []
        if(not useNetworkx):
            self.readDataFromTxt(path)
        else:
            if(networkxGraph == None):
                print("当你设置useNetworkx为True时，请指定networkxGraph字段")
            else:
                self.readDataFromNetworkx(networkxGraph)

    def cluster(self):
        # 网络聚类
        G = nx.Graph()
        G.add_nodes_from(range(len(self.nodeList)))

        for node in self.nodeList:
            nodeID = node.nodeID
            for adjacentNodeID in node.adjacentNodes:
                node1 = nodeID if nodeID < adjacentNodeID else adjacentNodeID
                node2 = nodeID if nodeID > adjacentNodeID else adjacentNodeID
                if not G.has_edge(node1, node2):
                    G.add_edge(node1, node2)
        
        initial_temp = 1.0      # 初始的温度
        cooling_rate = 0.99     # 冷却速度
        iterations = 10        # 迭代轮数
        min_community_size = 5 # 每个簇最小的节点个数
        max_community_size = 20 # 每个簇最大的节点个数

        partition = {node: i for i, node in enumerate(G.nodes)}
        communities = [[node] for node in G.nodes]  # 初始时每个节点自成一个社区
        modularity = nx.algorithms.community.modularity(G, communities)
        best_partition = partition
        best_modularity = modularity
        temp = initial_temp

        for _ in range(iterations):
            for node in G.nodes:
                # 统计当前社区的大小
                community_sizes = {com: len([n for n, c in partition.items() if c == com])
                                   for com in set(partition.values())}

                # 考虑只移动到已经达到最小尺寸的社区
                potential_communities = [com for com, size in community_sizes.items()
                                         if min_community_size <= size < max_community_size]

                # 如果当前社区大小小于最大值，可以考虑留在当前社区
                current_community = partition[node]
                if community_sizes[current_community] < max_community_size:
                    potential_communities.append(current_community)

                # 如果没有可行的社区，继续下一个节点
                if not potential_communities:
                    continue

                old_community = partition[node]

                # 尝试随机选择一个新的社区
                new_community = random.choice(potential_communities)
                partition[node] = new_community

                # 检测社区质量
                new_modularity = nx.algorithms.community.modularity(G, [frozenset([n for n, c in partition.items() if c == com])
                                                                       for com in set(partition.values())])

                delta_modularity = new_modularity - modularity
                # 如果增加了模块度或者冷却概率允许，则保留此次变动
                if delta_modularity > 0 or math.exp(delta_modularity / temp) > random.random():
                    modularity = new_modularity
                    if new_modularity > best_modularity:
                        best_partition = partition.copy()
                        best_modularity = new_modularity
                else:
                    partition[node] = old_community  # 撤回变更

            temp *= cooling_rate

        # 确保所有的社区都在指定的大小范围内，这可能需要一些后处理和分配策略

        # 输出最佳分区和模块度
        print("Best Partition:", best_partition)
        print("Best Modularity:", best_modularity)

    def getAllPathFromDFS(self, startNodeID, endNodeID, maxLength):
        def dfsHelper(currentNode, endNode, maxLength, currentPath, visited, paths):
            visited.add(currentNode)
            currentPath.append(currentNode)
            if currentNode == endNode:
                paths.append(list(currentPath))
            elif len(currentPath) <= maxLength:
                for adjacentNode in self.getNode(currentNode).adjacentNodes:
                    if adjacentNode not in visited:
                        dfsHelper(adjacentNode, endNode, maxLength, currentPath, visited, paths)
            currentPath.pop()
            visited.remove(currentNode)
        startNode = self.getNode(startNodeID).nodeID
        endNode = self.getNode(endNodeID).nodeID
        if startNode == -1 or endNode == -1:
            return [] 
        visited = set()
        paths = []
        dfsHelper(startNode, endNode, maxLength, [], visited, paths)
        return paths
        
    def getShortedPathWithOneBrokenLinkAndFlowMatrix(self, linkStartNodeID, linkEndNodeID, flowMatrix):
        distances = {}
        paths = {}
        visited = set()
        for node in self.nodeList:
            distances[node.nodeID] = float('-inf')  # 用负无穷初始化最大容量
            paths[node.nodeID] = []
        distances[linkStartNodeID] = float('inf')  # 起始节点的容量为正无穷
        while len(visited) < len(self.nodeList):
            max_capacity = float('-inf')
            max_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] > max_capacity:
                    max_capacity = distances[node.nodeID]
                    max_node = node
            if max_node is None:
                break
            visited.add(max_node.nodeID)
            for neighborID in max_node.adjacentNodes:
                # 跳过破损的链路
                if (max_node.nodeID == linkStartNodeID and neighborID == linkEndNodeID) or (max_node.nodeID == linkEndNodeID and neighborID == linkStartNodeID):
                    continue
                neighbor = self.getNode(neighborID)
                link_capacity = self.getLinkCapacity(max_node.nodeID, neighbor.nodeID)
                flow = flowMatrix[max_node.nodeID-1][neighbor.nodeID-1]
                available_capacity = link_capacity - flow  # 允许负数
                new_capacity = min(distances[max_node.nodeID], available_capacity)
                if new_capacity > distances[neighbor.nodeID]:
                    distances[neighbor.nodeID] = new_capacity
                    paths[neighbor.nodeID] = paths[max_node.nodeID] + [neighbor.nodeID]
        return [linkStartNodeID, *paths[linkEndNodeID]] if paths[linkEndNodeID] else []
        
    def getShortedPathWithOneBrokenLink(self, linkStartNodeID, linkEndNodeID):
        # 创建一个字典来保存节点的距离和路径
        distances = {}
        paths = {}
        visited = set()
        # 初始化距离为无穷大，路径为空列表
        for node in self.nodeList:
            distances[node.nodeID] = float('inf')
            paths[node.nodeID] = []
        # 将起始节点的距离设置为0
        distances[linkStartNodeID] = 0

        while len(visited) < len(self.nodeList):
            # 找到未访问节点中距离最小的节点
            min_distance = float('inf')
            min_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] < min_distance:
                    min_distance = distances[node.nodeID]
                    min_node = node
            if min_node is None:
                # 没有找到可到达的节点
                break
            # 标记节点为已访问
            visited.add(min_node.nodeID)
            # 更新相邻节点的距离和路径
            for neighborID in min_node.adjacentNodes:
                # 跳过破损的链路
                if (min_node.nodeID == linkStartNodeID and neighborID == linkEndNodeID) or (min_node.nodeID == linkEndNodeID and neighborID == linkStartNodeID):
                    continue
                neighbor = self.getNode(neighborID)
                # 计算新的距离
                new_distance = distances[min_node.nodeID] + 1
                if new_distance < distances[neighbor.nodeID]:
                    # 更新距离和路径
                    distances[neighbor.nodeID] = new_distance
                    paths[neighbor.nodeID] = paths[min_node.nodeID] + [neighbor.nodeID]

        return [linkStartNodeID, *paths[linkEndNodeID]] if distances[linkEndNodeID] != float('inf') else []

    # flowMatrix
    def getGreedyPathWithFlowMatrix(self, startNodeID, endNodeID, flowMatrix):
        # nodeNum = len(n.nodeList)
        # flowMatrix = [[0] * nodeNum for _ in range(nodeNum)]
        distances = {}
        paths = {}
        visited = set()
        for node in self.nodeList:
            distances[node.nodeID] = float('-inf')  # 用负无穷初始化最大容量
            paths[node.nodeID] = []
        distances[startNodeID] = float('inf')  # 起始节点的容量为正无穷
        while len(visited) < len(self.nodeList):
            max_capacity = float('-inf')
            max_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] > max_capacity:
                    max_capacity = distances[node.nodeID]
                    max_node = node
            if max_node is None:
                break
            visited.add(max_node.nodeID)
            for neighborID in max_node.adjacentNodes:
                neighbor = self.getNode(neighborID)
                link_capacity = self.getLinkCapacity(max_node.nodeID, neighbor.nodeID)
                flow = flowMatrix[max_node.nodeID-1][neighbor.nodeID-1]
                available_capacity = link_capacity - flow  # 允许负数
                new_capacity = min(distances[max_node.nodeID], available_capacity)
                if new_capacity > distances[neighbor.nodeID]:
                    distances[neighbor.nodeID] = new_capacity
                    paths[neighbor.nodeID] = paths[max_node.nodeID] + [neighbor.nodeID]
        return [startNodeID, *paths[endNodeID]] if paths[endNodeID] else []

    def getShortedPath(self, startNodeID, endNodeID):
        # 创建一个字典来保存节点的距离和路径
        distances = {}
        paths = {}
        visited = set()
        # 初始化距离为无穷大，路径为空列表
        for node in self.nodeList:
            distances[node.nodeID] = float('inf')
            paths[node.nodeID] = []
        # 将起始节点的距离设置为0
        distances[startNodeID] = 0
        while len(visited) < len(self.nodeList):
            # 找到未访问节点中距离最小的节点
            min_distance = float('inf')
            min_node = None
            for node in self.nodeList:
                if node.nodeID not in visited and distances[node.nodeID] < min_distance:
                    min_distance = distances[node.nodeID]
                    min_node = node
            if min_node is None:
                # 没有找到可到达的节点
                break
            # 标记节点为已访问
            visited.add(min_node.nodeID)
            # 更新相邻节点的距离和路径
            for neighborID in min_node.adjacentNodes:
                neighbor = self.getNode(neighborID)
                # 计算新的距离
                new_distance = distances[min_node.nodeID] + 1
                if new_distance < distances[neighbor.nodeID]:
                    # 更新距离和路径
                    distances[neighbor.nodeID] = new_distance
                    paths[neighbor.nodeID] = paths[min_node.nodeID] + [neighbor.nodeID]
        return paths[endNodeID]
        
    def showNetwork(self):
        with open('generate.txt', 'r') as file:
            lines = file.readlines()
        n = int(lines[0])
        nodes = [int(line) for line in lines[1:n+1]]
        G = nx.Graph()
        G.add_nodes_from(nodes)
        m = int(lines[n+1])
        for line in lines[n+2:]:
            start, end, weight = map(int, line.split())
            G.add_edge(start, end, weight=weight)
        fig, ax = plt.subplots(figsize=(10, 10))
        pos = nx.spring_layout(G, seed=41) 
        nx.draw(G, pos, ax=ax, with_labels=True, node_color='red', node_size=300, font_size=8, alpha=1)
        edge_labels = nx.get_edge_attributes(G, 'weight')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=6)
        plt.title("Network Topology")
        plt.show()
        # plt.savefig("topo.svg", format="svg")
    
    def getLinkNums(self):
        count = 0
        for node in self.nodeList:
            for _ in node.adjacentNodes:
                count += 1
        return count

    def getLinkCapacity(self, startID, endID):
        for linkCapacity in self.linkListCapacity:
            if linkCapacity.startNode == startID and linkCapacity.endNode == endID:
                return linkCapacity.capacity
            if linkCapacity.startNode == endID and linkCapacity.endNode == startID:
                return linkCapacity.capacity
        return -1
    
    def setLinkCapacity(self, startID, endID, capacity):
        self.linkListCapacity.append(self.LinkCapacity(startID, endID, capacity))

    def getNode(self, nodeId):
        if nodeId - 1 >= len(self.nodeList) or nodeId - 1 < 0:
            return self.Node(-1)
        return self.nodeList[nodeId - 1]

    def readDataFromNetworkx(self, graph: Graph):
        for nodeID in range(1, graph.number_of_nodes()+1):
            self.nodeList.append(self.Node(nodeID))
        for (u, v) in graph.edges():
            self.nodeList[u - 1].adjacentNodes.append(v)
            self.nodeList[v - 1].adjacentNodes.append(u)
            self.linkListCapacity.append(self.LinkCapacity(u, v, graph.get_edge_data(u, v)['weight']))

    def readDataFromTxt(self, path):
        with open(path, 'r') as file:
            lines = file.readlines()
            numNodes = int(lines[0])
            nodeIDs = [int(line.strip()) for line in lines[1 : numNodes + 1]]
            
            for nodeID in nodeIDs:
                self.nodeList.append(self.Node(nodeID))
            
            numLinks = int(lines[numNodes + 1])
            for i in range(numNodes + 2, numNodes + numLinks + 2):
                linkData = lines[i].split()
                startNode = int(linkData[0])
                endNode = int(linkData[1])
                self.nodeList[startNode - 1].adjacentNodes.append(endNode)
                self.nodeList[endNode - 1].adjacentNodes.append(startNode)
                capacity = int(linkData[2])
                linkCapacity = self.LinkCapacity(startNode, endNode, capacity)
                self.linkListCapacity.append(linkCapacity)
